import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras import utils

DATADIR = '../dl2_DataSets/zoo'

import cv2
import numpy as np

def readData(path):
    x_images = []
    y_labels = []
    for i, j in enumerate(os.listdir(path)):
        sub_path = os.path.join(path, j)
        for image_name in os.listdir(sub_path):
            image_path = os.path.join(sub_path, image_name)
            image = cv2.imread(image_path)/255
            image = cv2.resize(image, (100, 100))
            x_images.append(image)
            y_labels.append(i)
    return np.array(x_images), np.array(y_labels)

x_images, y_lables = readData(DATADIR)

y_lables = utils.to_categorical(y_lables, 2)
print("x_images, y_lables:", x_images.shape, y_lables.shape)

from sklearn.model_selection import train_test_split
x_train, x_test, y_train, y_test = train_test_split(x_images, y_lables, train_size=0.8, random_state=123)

# data pipeline
db_train = tf.data.Dataset.from_tensor_slices((x_train, y_train)).shuffle(100).batch(batch_size=8)
db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test)).shuffle(100).batch(batch_size=8)


class ConvCell(tf.keras.Model):
    def __init__(self, ch, strides, kernel_size=3, padding='same', scale=None):
        super(ConvCell, self).__init__()
        self.model = tf.keras.Sequential([
            tf.keras.layers.Conv2D(ch, kernel_size, strides=strides, padding=padding),
            tf.keras.layers.BatchNormalization(scale=scale)])
    def call(self, inputs):
        x = self.model(inputs, training=False)
        return x

class ResnetBlock(tf.keras.Model):
    def __init__(self, ch, strides):
        super(ResnetBlock, self).__init__()
        self.strides = strides
        self.ch = ch
        self.conv1 = ConvCell(self.ch, strides=strides)
        self.relu = tf.keras.layers.ReLU()
        self.conv2 = ConvCell(self.ch, strides=(1, 1))

        if strides != 1:
            self.residual = tf.keras.layers.Conv2D(self.ch, kernel_size=(1, 1), strides=strides)
        else:
            self.residual = lambda x: x

    def call(self, inputs):
        x = self.conv1(inputs)
        x = self.relu(x)
        x = self.conv2(x)
        r = self.residual(inputs)
        x = x + r
        x = self.relu(x)
        return x

class Resnet18(tf.keras.Model):
    def __init__(self, num_layers, num_class, ch=64):
        super(Resnet18, self).__init__()
        self.conv = ConvCell(ch, strides=(1, 1), scale=False)
        self.relu = tf.keras.layers.ReLU()
        self.out_ch = ch
        self.maxpool = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(1, 1), padding='same')

        self.blocks = tf.keras.Sequential()
        for block_id in range(num_layers):
            if block_id == 0:
                for i in range(2):
                    self.blocks.add(ResnetBlock(self.out_ch, strides=(1, 1)))
            else:
                self.blocks.add(ResnetBlock(self.out_ch, strides=(2, 2)))
                self.blocks.add(ResnetBlock(self.out_ch, strides=(1, 1)))
            self.out_ch *= 2
        self.gapool = tf.keras.layers.GlobalAveragePooling2D()
        self.fc = tf.keras.layers.Dense(num_class, activation='softmax')

    def call(self, inputs):
        x = self.conv(inputs)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.blocks(x)
        x = self.gapool(x)
        x = self.fc(x)
        return x


if __name__ == '__main__':

    model = Resnet18(num_layers=4, num_class=2)

    model.compile(loss=tf.keras.losses.CategoricalCrossentropy(),
                  optimizer=tf.keras.optimizers.Adam(1e-4),
                  metrics=['accuracy'])

    history = model.fit(db_train, validation_data=db_test, epochs=6)

    acc = history.history['accuracy']
    val_acc = history.history['val_accuracy']

    loss = history.history['loss']
    val_loss = history.history['val_loss']

    import matplotlib.pyplot as plt

    plt.plot(acc, label='acc')
    plt.plot(val_acc, label='val_acc')
    plt.legend()
    plt.show()

    plt.plot(loss, label='loss')
    plt.plot(val_loss, label='val_loss')
    plt.legend()
    plt.show()
